import numpy as np
import argparse

def close_ratio(gold, x, rtol=1e-2, eps=1e-3, eps_atol=1e-5, etol=1e-2):
    abs_gold = np.abs(gold)
    # mask
    mask = abs_gold <= eps
    # absolute error
    ae = np.abs(x - gold)
    # relative error
    re = ae / abs_gold
    re_pass = re <= rtol
    re_pass[mask] = True
    ae_pass = ae <= eps_atol
    ae_pass[np.logical_not(mask)] = True
    all_pass = np.logical_and(re_pass, ae_pass)
    pass_ratio = np.sum(all_pass) / np.size(all_pass)
    res = (1 - pass_ratio) <= etol
    print(f"Close ratio: {pass_ratio}, all_close: {res}")
    return pass_ratio, res

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--ms', type=str, required=True)
    parser.add_argument('--pt', type=str, required=True)
    parser.add_argument('--rtol', type=float, default=1e-3)
    parser.add_argument('--etol', type=float, default=1e-3)
    args = parser.parse_args()
    ms_npy = np.load(args.ms)
    pt_npy = np.load(args.pt)
    breakpoint()
    close_ratio(pt_npy, ms_npy, rtol=args.rtol, etol=args.etol)
